#include "Renderer.h"
#include <optix_function_table_definition.h> // REVIEW: i dont know why, but it worked ???

using namespace OBR;

Renderer::Renderer()
{
    this->initOptix();
}

Renderer::~Renderer()
{
    this->freeResources();
}

static void context_log_cb(unsigned int level, const char *tag, const char *message,
                           void * /*cbdata */)
{
    std::cerr << "[" << std::setw(2) << level << "][" << std::setw(12) << tag << "]: " << message
              << "\n";
}

void Renderer::initOptix()
{
    // init cuda and optix
    CUDA_CHECK(cudaFree(0));
    CUcontext cu_ctx = 0;
    OPTIX_CHECK(optixInit());

    // write options
    OptixDeviceContextOptions options = {};
    options.logCallbackFunction = &context_log_cb;
    options.logCallbackLevel = 4;

    // create context
    OPTIX_CHECK(optixDeviceContextCreate(cu_ctx, &options, &(this->context)));

    // create cuda stream
    CUDA_CHECK(cudaStreamCreate(&(this->stream)));

    // change device settings
    size_t heap_size = 4 * 1024 * 1024 * 1024; // 4 GB
    CUDA_CHECK(cudaDeviceSetLimit(cudaLimitMallocHeapSize, heap_size));
}

void Renderer::freeResources()
{
    CUDA_CHECK(cudaFree(0));
}

void Renderer::createPipeline(OptixPipelineCompileOptions &pipeline_compile_options)
{
    // TODO: use configurator settings
    const uint32_t max_trace_depth = 8;
    const uint32_t max_cc_depth = 2;
    const uint32_t max_dc_depth = 2;
    const uint32_t max_traversal_depth = 4;

    // Link program groups to pipeline
    OptixPipelineLinkOptions pipeline_link_options = {
        max_trace_depth,                  // maxTraceDepth
        OPTIX_COMPILE_DEBUG_LEVEL_DEFAULT // debugLevel
    };
    char *log = new char[MAX_INTERNAL_LOG_SIZE];
    size_t log_size = MAX_INTERNAL_LOG_SIZE;
    OPTIX_CHECK(optixPipelineCreate(this->context, &pipeline_compile_options,
                                    &pipeline_link_options, CudaProgram::all_program_groups.data(),
                                    CudaProgram::all_program_groups.size(), log, &log_size,
                                    &(this->pipeline));
                std::cout << "=============== Pipeline Create Log ===============" << std::endl;
                std::cout << log << std::endl; delete[] log;);

    // calculate and set stack sizes
    OptixStackSizes stack_sizes = {};
    for (OptixProgramGroup &prog_group : CudaProgram::all_program_groups)
    {
        OPTIX_CHECK(optixUtilAccumulateStackSizes(prog_group, &stack_sizes));
    }
    uint32_t direct_callable_stack_size_from_traversal;
    uint32_t direct_callable_stack_size_from_state;
    uint32_t continuation_stack_size;
    OPTIX_CHECK(optixUtilComputeStackSizes(&stack_sizes, max_trace_depth, max_cc_depth,
                                           max_dc_depth, &direct_callable_stack_size_from_traversal,
                                           &direct_callable_stack_size_from_state,
                                           &continuation_stack_size));
    OPTIX_CHECK(optixPipelineSetStackSize(this->pipeline, direct_callable_stack_size_from_traversal,
                                          direct_callable_stack_size_from_state,
                                          continuation_stack_size, max_traversal_depth));
}

void Renderer::setRaygenSbt(std::shared_ptr<Camera> &cam, uint med_id_offset)
{
    // update and set sbt info
    cam->uploadSbtData();
    RaygenRecord record;
    OPTIX_CHECK(optixSbtRecordPackHeader(cam->getProgram()->program_group, &record));
    record.data.camera_payload = cam->getSbtDevicePtr();
    record.data.med_id = cam->med_id;
    record.data.med_id += cam->med_id != -1 ? med_id_offset : 0;
    record.data.med_instance_id = cam->med_instance_id;
    d_raygen_sbt_record.resize(sizeof(RaygenRecord));
    d_raygen_sbt_record.upload(&record, sizeof(RaygenRecord));
    this->sbt.raygenRecord = d_raygen_sbt_record.d_pointer();
}

void Renderer::setMissSbt(std::shared_ptr<Background> &bg)
{
    // update and set sbt info
    bg->uploadSbtData();
    MissRecord record;
    OPTIX_CHECK(optixSbtRecordPackHeader(bg->getProgram()->program_group, &record));
    record.data.background_payload = bg->getSbtDevicePtr();
    d_miss_sbt_record.resize(sizeof(MissRecord));
    d_miss_sbt_record.upload(&record, sizeof(MissRecord));
    this->sbt.missRecordBase = d_miss_sbt_record.d_pointer();
    this->sbt.missRecordStrideInBytes = sizeof(MissRecord);
    this->sbt.missRecordCount = 1u;
}

void Renderer::setHitSbt(uint mat_id_offset, uint med_id_offset,
                         VirtualProgramManager &mesh_manager)
{
    // init constants
    const size_t num_records = StaticInstance::hit_data.size();
    const size_t record_size = sizeof(HitRecord);
    const size_t memory_size = num_records * record_size;

    // create records
    HitRecord *tmp_records = new HitRecord[num_records];
    for (size_t i = 0; i < num_records; i++)
    {
        std::shared_ptr<Mesh> mesh = StaticInstance::hit_meshes[i];
        if (mesh->gas_sbt_base != -1)
        {
            size_t sbt_idx = mesh->gas_sbt_base + StaticInstance::hit_gas_offset[i];
            OPTIX_CHECK(
                optixSbtRecordPackHeader(mesh->getProgram()->program_group, tmp_records + sbt_idx));
            HitData &data = tmp_records[sbt_idx].data;
            data = StaticInstance::hit_data[i];
            // data.mat_id += data.mat_id >= 0 ? mat_id_offset : 0;
            // data.med_id_inside += data.med_id_inside >= 0 ? med_id_offset : 0;
            // data.med_id_outside += data.med_id_outside >= 0 ? med_id_offset : 0;
            data.hit_payload = mesh->getPayloadDevicePtr();
            data.n_instances = mesh_manager.getInstanceCount(typeid(*mesh));
        }
    }

    // migrate memory and release temp buffer
    d_hit_sbt_records.resize(memory_size);
    d_hit_sbt_records.upload(tmp_records, memory_size);
    delete[] tmp_records;

    // set sbt info
    this->sbt.hitgroupRecordBase = d_hit_sbt_records.d_pointer();
    this->sbt.hitgroupRecordCount = num_records;
    this->sbt.hitgroupRecordStrideInBytes = record_size;
}

void Renderer::setCallableSbt(std::shared_ptr<Integrator> &integrator,
                              VirtualProgramManager &light_manager,
                              VirtualProgramManager &material_manager,
                              VirtualProgramManager &medium_manager)
{
    const size_t n_mediums = medium_manager.getUniqueCount();
    const size_t n_materials = material_manager.getUniqueCount();
    const size_t n_lights = light_manager.getUniqueCount();
    const size_t n_records = 1 + n_mediums + n_materials + n_lights;
    const size_t record_size = sizeof(CallableRecord);
    const size_t memory_size = n_records * record_size;

    // init callable records
    CallableRecord *tmp_records = new CallableRecord[n_records];
    // integrator
    integrator->uploadSbtData();
    OptixProgramGroup &prog = integrator->getProgram()->program_group;
    OPTIX_CHECK(optixSbtRecordPackHeader(prog, tmp_records));
    tmp_records[0].data.callable_payload = integrator->getSbtDevicePtr();
    // others
    size_t i = 1;
    for (auto &&c : light_manager.getUnique())
    {
        OptixProgramGroup &prog = c->getProgram()->program_group;
        OPTIX_CHECK(optixSbtRecordPackHeader(prog, tmp_records + i));
        tmp_records[i].data.callable_payload = c->getPayloadDevicePtr();
        tmp_records[i].data.n_instances = light_manager.getInstanceCount(typeid(*c));
        i++;
    }
    for (auto &&c : material_manager.getUnique())
    {
        OptixProgramGroup &prog = c->getProgram()->program_group;
        OPTIX_CHECK(optixSbtRecordPackHeader(prog, tmp_records + i));
        tmp_records[i].data.callable_payload = c->getPayloadDevicePtr();
        tmp_records[i].data.n_instances = material_manager.getInstanceCount(typeid(*c));
        i++;
    }
    for (auto &&c : medium_manager.getUnique())
    {
        OptixProgramGroup &prog = c->getProgram()->program_group;
        OPTIX_CHECK(optixSbtRecordPackHeader(prog, tmp_records + i));
        tmp_records[i].data.callable_payload = c->getPayloadDevicePtr();
        tmp_records[i].data.n_instances = medium_manager.getInstanceCount(typeid(*c));
        i++;
    }

    // migrate memory and release temp buffer
    d_callable_sbt_records.resize(memory_size);
    d_callable_sbt_records.upload(tmp_records, memory_size);
    delete[] tmp_records;

    // set sbt info
    this->sbt.callablesRecordBase = d_callable_sbt_records.d_pointer();
    this->sbt.callablesRecordCount = n_records;
    this->sbt.callablesRecordStrideInBytes = record_size;
}

void Renderer::createSbt(std::shared_ptr<Camera> &camera,         // raygen
                         std::shared_ptr<Background> &background, // miss,
                         VirtualProgramManager &mesh_manager,     // hit
                         std::shared_ptr<Integrator> &integrator, // callable
                         VirtualProgramManager &light_manager,    // callable
                         VirtualProgramManager &material_manager, // callable
                         VirtualProgramManager &medium_manager    // callable
)
{
    const size_t n_materials = material_manager.getUniqueCount();
    const size_t n_lights = light_manager.getUniqueCount();
    const size_t mat_id_offset = 1 + n_lights;
    const size_t med_id_offset = mat_id_offset + n_materials;

    this->sbt = OptixShaderBindingTable();
    this->setRaygenSbt(camera, med_id_offset);
    this->setMissSbt(background);
    this->setHitSbt(mat_id_offset, med_id_offset, mesh_manager);
    this->setCallableSbt(integrator, light_manager, material_manager, medium_manager);
    this->sbt.exceptionRecord = 0u;
}

void Renderer::renderFrame(Params &params)
{
    float dt; // delta time
    auto t0 = std::chrono::steady_clock::now();

    this->prepareLaunchParams(params);

    auto t1 = std::chrono::steady_clock::now();
    dt = std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count() / 1e9;
    // printf("time used in setting up: %.6f sec(s)\n", dt);

    OPTIX_CHECK(optixLaunch(this->pipeline, this->stream,                      // contex & pipeline
                            this->d_launch_params.d_pointer(), sizeof(Params), // params
                            &(this->sbt),                                      // sbt
                            params.width,                                      // launch width
                            params.height,                                     // launch height
                            1u                                                 // launch depth
                            ));
    // synchronize
    CUDA_CHECK(cudaStreamSynchronize(this->stream));
    CUDA_CHECK(cudaDeviceSynchronize());

    auto t2 = std::chrono::steady_clock::now();
    dt = std::chrono::duration_cast<std::chrono::nanoseconds>(t2 - t1).count() / 1e9;
    // printf("time used in rendering one frame: %.6f sec(s)\n", dt);

    this->downloadResult(params);

    auto t3 = std::chrono::steady_clock::now();
    dt = std::chrono::duration_cast<std::chrono::nanoseconds>(t3 - t2).count() / 1e9;
    // printf("time used in downloading result: %.6f sec(s)\n", dt);
}

void Renderer::prepareLaunchParams(Params &base)
{
    Params launch_params(base);
    launch_params.subframe_index = 0;

    // set time
    launch_params.launch_time = std::chrono::duration_cast<std::chrono::milliseconds>(
                                    std::chrono::system_clock::now().time_since_epoch())
                                    .count();

    // resize fbo
    size_t fbo_size = sizeof(Spectrum) * base.width * base.height * base.num_fbos;
    this->d_fbos.resize(fbo_size);

    // bind pre allocated memory
    launch_params.fbo = reinterpret_cast<Spectrum *>(this->d_fbos.d_pointer());

    // upload launch params
    this->d_launch_params.resize(sizeof(Params));
    this->d_launch_params.upload(&launch_params, sizeof(Params));

    // compute and log total buffer size
    // size_t total_buffer_size = fbo_size;
    // float size_in_MiB = total_buffer_size / (float)(1 << 20);
    // printf("total buffer size: %.6f MiB\n", size_in_MiB);
}

void Renderer::downloadResult(Params &dest)
{
    if (dest.fbo != nullptr)
        delete[] dest.fbo;

    // download fbos
    size_t fbo_size = dest.width * dest.height * dest.num_fbos;
    dest.fbo = new Spectrum[fbo_size];
    this->d_fbos.download(dest.fbo, fbo_size * sizeof(Spectrum));

    // download trackers

    // download pbos
}